package main

import (
	"context"
	"fmt"
	"log"
	"strings"
	"time"

	"mcp-server/internal/config"
	"mcp-server/internal/storage"
	"mcp-server/internal/vector"
)

type SearchDemo struct {
	storageManager *storage.Manager
	vectorService  *vector.Service
}

func NewSearchDemo() (*SearchDemo, error) {
	// 初始化配置
	cfg := &config.Config{
		Storage: config.StorageConfig{
			Type: "simple",
			Qdrant: config.QdrantConfig{
				Host:     "localhost",
				Port:     6333,
				DataPath: "./knowledge_vectors",
			},
		},
		Vector: config.VectorConfig{
			Model:     "bge-base-zh-v1.5",
			Dimension: 768,
			BatchSize: 32,
		},
	}

	// 初始化存储管理器
	log.Println("正在初始化存储管理器...")
	storageManager, err := storage.NewManager(cfg)
	if err != nil {
		return nil, fmt.Errorf("初始化存储管理器失败: %v", err)
	}

	// 初始化向量化服务
	log.Println("正在初始化向量化服务...")
	vectorService, err := vector.NewService(cfg)
	if err != nil {
		return nil, fmt.Errorf("初始化向量化服务失败: %v", err)
	}

	return &SearchDemo{
		storageManager: storageManager,
		vectorService:  vectorService,
	}, nil
}

func (sd *SearchDemo) Close() {
	log.Println("正在关闭搜索演示...")
	if sd.vectorService != nil {
		sd.vectorService.Close()
	}
	if sd.storageManager != nil {
		sd.storageManager.Close()
	}
}

func (sd *SearchDemo) searchAndDisplay(query string, topK int) error {
	fmt.Printf("\n=== 搜索查询: \"%s\" ===\n", query)

	ctx := context.Background()
	start := time.Now()

	// 执行搜索
	results, err := sd.storageManager.SearchVectors(ctx, "knowledge_base", query, topK)
	if err != nil {
		return fmt.Errorf("搜索失败: %v", err)
	}

	duration := time.Since(start)
	fmt.Printf("搜索耗时: %v\n", duration)
	fmt.Printf("找到 %d 个相关结果:\n\n", len(results))

	// 显示搜索结果
	for i, result := range results {
		fmt.Printf("结果 %d (相似度: %.4f):\n", i+1, result.Score)
		fmt.Printf("  文档ID: %s\n", result.ID)

		// 解析文件名
		parts := strings.Split(result.ID, "_chunk_")
		if len(parts) > 0 {
			fileName := extractFileName(parts[0])
			fmt.Printf("  文件名: %s\n", fileName)
		}

		// 显示文本内容
		if result.Metadata != nil {
			if text, ok := result.Metadata["text"].(string); ok {
				// 限制显示长度
				displayText := text
				if len(displayText) > 200 {
					displayText = displayText[:200] + "..."
				}
				fmt.Printf("  内容: %s\n", displayText)
			}
		}
		fmt.Println()
	}

	return nil
}

func extractFileName(path string) string {
	// 从路径中提取文件名
	parts := strings.Split(path, "_")
	if len(parts) > 0 {
		// 查找.md部分
		for i, part := range parts {
			if strings.Contains(part, ".md") {
				// 重新组合文件名
				fileName := strings.Join(parts[len(parts)-1:], "_")
				if strings.Contains(fileName, ".md") {
					return fileName
				}
				// 如果没找到，尝试其他方式
				for j := i; j < len(parts); j++ {
					candidate := strings.Join(parts[j:], "_")
					if strings.Contains(candidate, ".md") {
						return candidate
					}
				}
			}
		}
	}
	return path
}

func (sd *SearchDemo) runDemoQueries() error {
	// 定义测试查询
	queries := []struct {
		query string
		topK  int
		desc  string
	}{
		{"国有企业监管", 5, "基础监管查询"},
		{"财务预算管理", 3, "财务管理相关"},
		{"党建工作", 3, "党建职责相关"},
		{"风险管理", 3, "风险控制相关"},
		{"资产评估", 3, "资产管理相关"},
		{"薪酬管理", 3, "薪酬制度相关"},
		{"投资决策", 3, "投资管理相关"},
		{"责任追究", 3, "责任制度相关"},
	}

	fmt.Println("=== 知识库搜索演示 ===")
	fmt.Printf("知识库包含 %d 个文档的向量化内容\n", 15)

	for _, q := range queries {
		fmt.Printf("\n--- %s ---\n", q.desc)
		if err := sd.searchAndDisplay(q.query, q.topK); err != nil {
			log.Printf("查询 '%s' 失败: %v", q.query, err)
		}
		time.Sleep(500 * time.Millisecond) // 短暂延迟
	}

	return nil
}

func (sd *SearchDemo) getCollectionStats() error {
	fmt.Println("\n=== 知识库统计信息 ===")

	ctx := context.Background()

	// 检查集合是否存在
	exists, err := sd.storageManager.CollectionExists(ctx, "knowledge_base")
	if err != nil {
		return fmt.Errorf("检查集合失败: %v", err)
	}

	if !exists {
		fmt.Println("知识库集合不存在")
		return nil
	}

	// 获取集合信息
	info, err := sd.storageManager.GetCollectionInfo(ctx, "knowledge_base")
	if err != nil {
		return fmt.Errorf("获取集合信息失败: %v", err)
	}

	fmt.Printf("集合名称: %s\n", info.Name)
	fmt.Printf("向量数量: %d\n", info.VectorCount)
	fmt.Printf("向量维度: %d\n", info.VectorDimension)
	fmt.Printf("距离度量: %s\n", info.DistanceMetric)
	fmt.Printf("集合状态: %s\n", info.Status)

	return nil
}

func main() {
	// 创建搜索演示实例
	demo, err := NewSearchDemo()
	if err != nil {
		log.Fatalf("初始化搜索演示失败: %v", err)
	}
	defer demo.Close()

	// 显示统计信息
	if err := demo.getCollectionStats(); err != nil {
		log.Printf("获取统计信息失败: %v", err)
	}

	// 运行演示查询
	if err := demo.runDemoQueries(); err != nil {
		log.Fatalf("运行演示查询失败: %v", err)
	}

	// 交互式搜索
	fmt.Println("\n=== 交互式搜索 ===")
	fmt.Println("输入查询文本进行搜索 (输入 'quit' 退出):")

	for {
		fmt.Print("\n请输入查询: ")
		var query string
		if _, err := fmt.Scanln(&query); err != nil {
			continue
		}

		if query == "quit" || query == "exit" {
			break
		}

		if strings.TrimSpace(query) == "" {
			continue
		}

		if err := demo.searchAndDisplay(query, 5); err != nil {
			log.Printf("搜索失败: %v", err)
		}
	}

	fmt.Println("搜索演示结束")
}
